package cn.yx.common.mdc.spring;

import cn.yx.common.mdc.wrapper.MdcWrapper;
import org.slf4j.MDC;
import org.springframework.boot.autoconfigure.task.TaskExecutionProperties;
import org.springframework.lang.NonNull;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.time.Duration;
import java.util.Objects;
import java.util.concurrent.Callable;
import java.util.concurrent.Future;

/**
 * <p>MDC装饰Spring ThreadPoolTaskExecutor，可支持MDC上下文信息</p>
 *
 * @author Wgssmart
 */
public class MdcThreadPoolTaskExecutor extends ThreadPoolTaskExecutor {

    public MdcThreadPoolTaskExecutor(TaskExecutionProperties taskExecutionProperties) {
        setCorePoolSize(taskExecutionProperties.getPool().getCoreSize());
        setMaxPoolSize(taskExecutionProperties.getPool().getMaxSize());
        setQueueCapacity(taskExecutionProperties.getPool().getQueueCapacity());
        setKeepAliveSeconds((int) taskExecutionProperties.getPool().getKeepAlive().getSeconds());
        setTaskDecorator(new MdcTaskDecorator());
        // shutdown时是否等待异步任务完成
        setWaitForTasksToCompleteOnShutdown(taskExecutionProperties.getShutdown().isAwaitTermination());
        // 等待时间 （默认为0，此时立即停止），并没等待xx秒后强制停止
        Duration awaitTerminationPeriod = taskExecutionProperties.getShutdown().getAwaitTerminationPeriod();
        if (Objects.nonNull(awaitTerminationPeriod)) {
            setAwaitTerminationSeconds((int) taskExecutionProperties.getShutdown().getAwaitTerminationPeriod().getSeconds());
        }
        setThreadNamePrefix(taskExecutionProperties.getThreadNamePrefix());
        initialize();
    }

    @Override
    public void execute(@NonNull Runnable command) {
        super.execute(MdcWrapper.wrap(command, MDC.getCopyOfContextMap()));
    }

    @Override
    @NonNull
    public Future<?> submit(@NonNull Runnable task) {
        return super.submit(MdcWrapper.wrap(task, MDC.getCopyOfContextMap()));
    }

    @Override
    @NonNull
    public <T> Future<T> submit(@NonNull Callable<T> task) {
        return super.submit(MdcWrapper.wrap(task, MDC.getCopyOfContextMap()));
    }

}
